import random
from typing import List, Dict
from PIL import Image, ImageDraw
import numpy as np
import matplotlib.pyplot as plt
class RandomDotGenerator:
    def __init__(
        self,
        num_range=(1, 5),
        size_range=(0.01, 0.05),
        color_list=['red', 'orange', 'yellow', 'green', 'cyan', 'blue', 'purple'],
    ):
        self.num_range = num_range
        self.size_range = size_range
        self.color_list = color_list

    def generate(self, image: Image.Image, style: int = 0):
        """
        style:
            0 - 小圆点 (默认)
            1 - 小方块
            2 - 大圆点
            3 - 大方块
        """
        img_w, img_h = image.size
        dots_info = []
        draw = ImageDraw.Draw(image)

        large_shape = style in (2, 3)
        shape = 'circle' if style in (0, 2) else 'square'

        num_dots = 1 if large_shape else random.randint(*self.num_range)

        for _ in range(num_dots):
            size_pct = random.uniform(0.2, 0.4) if large_shape else random.uniform(*self.size_range)
            size_pixel = size_pct * min(img_w, img_h)

            position_pct = [
                round(random.uniform(0.0, 1.0), 4),
                round(random.uniform(0.0, 1.0), 4)
            ]
            color = random.choice(self.color_list)

            cx, cy = position_pct[0] * img_w, position_pct[1] * img_h
            bbox = [cx - size_pixel / 2, cy - size_pixel / 2, cx + size_pixel / 2, cy + size_pixel / 2]

            if shape == 'circle':
                draw.ellipse(bbox, fill=color)
            else:
                draw.rectangle(bbox, fill=color)

            dot = {
                'position_pct': position_pct,
                'size_pct': round(size_pct, 4),
                'color': color,
                'shape': shape,
                'large_shape': large_shape
            }
            dots_info.append(dot)

        return image, {
            'image_size': [img_w, img_h],
            'num_dots': num_dots,
            'dots': dots_info,
            'large_shape': large_shape
        }

def get_grid_token(pos_pct, num_token=14):
    grid_x = min(int(pos_pct[0] * num_token), num_token - 1)
    grid_y = min(int(pos_pct[1] * num_token), num_token - 1)
    return grid_x, grid_y

def get_covered_grids(dot, num_token=14):
    pos_x, pos_y = dot['position_pct']
    size = dot['size_pct']

    half_size = size / 2
    x_min = max(pos_x - half_size, 0)
    x_max = min(pos_x + half_size, 1)
    y_min = max(pos_y - half_size, 0)
    y_max = min(pos_y + half_size, 1)

    grid_x_min = min(int(x_min * num_token), num_token - 1)
    grid_x_max = min(int(x_max * num_token), num_token - 1)
    grid_y_min = min(int(y_min * num_token), num_token - 1)
    grid_y_max = min(int(y_max * num_token), num_token - 1)

    grids = []
    for gx in range(grid_x_min, grid_x_max + 1):
        for gy in range(grid_y_min, grid_y_max + 1):
            grids.append(f"<grid_{gx}><grid_{gy}>")
    return grids

def construct_description(shapes_info, num_token=14, question_type=None):
    color_map = {
        'red': '红色', 'orange': '橙色', 'yellow': '黄色',
        'green': '绿色', 'cyan': '青色', 'blue': '蓝色', 'purple': '紫色'
    }

    if not shapes_info:
        raise ValueError(f"shape info should not be none")

    question, answer = "", ""

    if question_type == "位置识别":
        q_templates = [
            "图中有哪些标记？请说出它们所在的位置与颜色。",
            "请列出所有标记点的grid坐标与颜色。",
            "所有图中标记的坐标及颜色分别是什么？"
        ]
        a_lines = []
        for shape in shapes_info:
            gx, gy = get_grid_token(shape['position_pct'], num_token)
            a_lines.append(f"<grid_{gx}><grid_{gy}>，颜色{color_map[shape['color']]}")
        question = random.choice(q_templates)
        answer = "；".join(a_lines) + "。"

    elif question_type == "邻近判断":
        ref_grid = random.randint(1, num_token - 2), random.randint(1, num_token - 2)
        q_templates = [
            f"请问<grid_{ref_grid[0]}><grid_{ref_grid[1]}>周围九宫格内是否存在标记？",
            f"以<grid_{ref_grid[0]}><grid_{ref_grid[1]}>为中心的九格内有标记点吗？"
        ]
        found = False
        a_lines = []
        for shape in shapes_info:
            gx, gy = get_grid_token(shape['position_pct'], num_token)
            if abs(gx - ref_grid[0]) <= 1 and abs(gy - ref_grid[1]) <= 1:
                found = True
                a_lines.append(f"有标记在<grid_{gx}><grid_{gy}>，颜色{color_map[shape['color']]}")
        question = random.choice(q_templates)
        answer = "；".join(a_lines) + "。" if found else "该九宫格范围内没有标记点。"

    elif question_type == "方位推断":
        ref_grid = random.randint(0, num_token - 1), random.randint(0, num_token - 1)
        directions = {'左': (-1, 0), '右': (1, 0), '上': (0, -1), '下': (0, 1)}
        direction = random.choice(list(directions.keys()))
        dx, dy = directions[direction]
        q_templates = [
            f"<grid_{ref_grid[0]}><grid_{ref_grid[1]}>的{direction}侧相邻区域是哪一个？",
            f"如果我在<grid_{ref_grid[0]}><grid_{ref_grid[1]}>，那么{direction}边是哪一格？"
        ]
        nx, ny = ref_grid[0] + dx, ref_grid[1] + dy
        if not (0 <= nx < num_token and 0 <= ny < num_token):
            answer = f"已经处于{direction}边界，无法继续移动。"
        else:
            answer = f"{direction}边的区域为<grid_{nx}><grid_{ny}>。"
        question = random.choice(q_templates)

    elif question_type == "颜色形状综合":
        shape = random.choice(shapes_info)
        shape_name = "圆形" if shape['shape'] == 'circle' else "方形"
        color = color_map[shape['color']]
        gx, gy = get_grid_token(shape['position_pct'], num_token)
        q_templates = [
            f"图中是否存在{color}的{shape_name}标记？如果有，它在哪？",
            f"请指出{color}的{shape_name}在grid中的位置。"
        ]
        question = random.choice(q_templates)
        answer = f"存在，位于<grid_{gx}><grid_{gy}>。"

    elif question_type == "大区域覆盖":
        large_shapes = [s for s in shapes_info if s['large_shape']]
        q_templates = [
            "请问图中的大形状覆盖了哪些区域？",
            "图中大标记的grid范围是哪些？"
        ]
        if large_shapes:
            shape = random.choice(large_shapes)
            shape_name = "大圆形" if shape['shape'] == 'circle' else "大方形"
            grids = get_covered_grids(shape, num_token)
            question = random.choice(q_templates)
            answer = f"{shape_name}标记覆盖区域包括：" + "、".join(grids) + "。"
        else:
            question = random.choice(q_templates)
            answer = "图中没有大区域标记。"

    return question.strip(), answer.strip()

if __name__ == "__main__":
    # Example usage
    # 参数设置
    num_token = 14  # 网格尺寸
    img_size = (224, 224)  # 图片大小
    generator = RandomDotGenerator()  # 点生成器
    
    # 模拟不同样式的问题类型
    styles = [0, 1, 2, 3]  # 假设4种风格：0, 1, 2, 3
    question_types = ["位置识别", "邻近判断", "方位推断", "颜色形状综合", "大区域覆盖"]

    # 可视化与调试
    for style in styles:
        print(f"样式 {style} 的问题类型：")
        
        # 生成图像并为其生成标记
        img = Image.fromarray(np.zeros((img_size[1], img_size[0], 3), dtype=np.uint8))  # 黑色背景图像
        img_with_dots, info = generator.generate(img, style=style)
        
        # 随机选择两个问题类型
        selected_question_types = random.sample(question_types, 2)
        
        # 为选中的问题类型生成问题
        for q_type in selected_question_types:
            txt = construct_description(info['dots'], num_token=num_token, question_type=q_type)
            print(f"问题类型: {q_type}")
            print(f"问题: {txt}")
            print("-" * 50)

        # 可视化生成的图像
        plt.imshow(img_with_dots)
        plt.title(f"样式 {style}")
        plt.axis('off')  # 关闭坐标轴
        plt.savefig(f"./experiments/test_grid/style_{style}_{q_type}.png")